package com.study.crypto.gb.server.advice;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.parser.Feature;
import com.study.crypto.dto.SignInfoDto;
import com.study.crypto.dto.gb.base.ResponseGbDto;
import com.study.crypto.gb.server.config.EnvironmentConfig;
import com.study.crypto.gb.server.dto.ReturnResult;
import com.study.crypto.gb.server.exception.SignatureException;
import com.study.crypto.signer.Signer;
import com.study.crypto.signer.SignerFactory;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.codec.binary.Base64;
import org.bouncycastle.asn1.ASN1InputStream;
import org.bouncycastle.asn1.ASN1Integer;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.gm.GMObjectIdentifiers;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.crypto.CryptoException;
import org.bouncycastle.pqc.legacy.math.linearalgebra.ByteUtils;
import org.bouncycastle.util.BigIntegers;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.MethodParameter;
import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.core.annotation.Order;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.bind.annotation.ResponseBody;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import java.io.IOException;
import java.security.GeneralSecurityException;
import java.security.PrivateKey;

/**
 * 对响应数据进行签名，方法上使用了 @RequireSignature 注解
 * @author Songjin
 * @since 2021-06-12 22:21
 */
@Slf4j
@Order(1)
@ControllerAdvice
public class ResponseBodySignatureAdvice implements ResponseBodyAdvice<Object> {
    
    private final Signer signer = SignerFactory.produce(GMObjectIdentifiers.sm2sign_with_sm3);
    @Autowired
    private EnvironmentConfig config;
    
    @Override
    public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
        return (AnnotatedElementUtils.hasAnnotation(returnType.getContainingClass(), ResponseBody.class) ||
                returnType.hasMethodAnnotation(ResponseBody.class));
    }
    
    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType, Class<? extends HttpMessageConverter<?>> selectedConverterType, ServerHttpRequest httpRequest, ServerHttpResponse httpResponse) {
        log.info("ResponseBodySignatureAdvice.beforeBodyWrite--->begin");
        boolean instanceof1 = body instanceof ReturnResult;
        if (!instanceof1) {
            log.info("ResponseBodySignatureAdvice.beforeBodyWrite--->end");
            return body;
        }
        Object data = ((ReturnResult) body).getData();
        boolean instanceof2 = data instanceof ResponseGbDto;
        if (!instanceof2) {
            log.info("ResponseBodySignatureAdvice.beforeBodyWrite--->end");
            return body;
        }
        ResponseGbDto response = (ResponseGbDto) data;
        PrivateKey privateKey = config.getPrivateKey();
        byte[] signature;
        try {
            JSONObject json = JSON.parseObject(JSON.toJSONString(response), Feature.OrderedField);
            String dataText = json.getString("data");
            json.put("data", dataText);
            byte[] inData = JSON.toJSONBytes(json);
            byte[] signature_ = signer.sign(inData, privateKey);
            try (ASN1InputStream inputStream = new ASN1InputStream(signature_)) {
                ASN1Sequence sequence = (ASN1Sequence) inputStream.readObject();
                ASN1Integer r_ = (ASN1Integer) sequence.getObjectAt(0);
                ASN1Integer s_ = (ASN1Integer) sequence.getObjectAt(1);
                byte[] r = BigIntegers.asUnsignedByteArray(r_.getPositiveValue());
                byte[] s = BigIntegers.asUnsignedByteArray(s_.getPositiveValue());
                signature = ByteUtils.concatenate(r, s);
            }
            log.info("返回业务参数签名值: {}", Base64.encodeBase64String(signature));
        } catch (IOException | CMSException | CryptoException | GeneralSecurityException e) {
            log.error("签名异常", e);
            throw new SignatureException("签名异常", e);
        }
        response.setSignInfo(new SignInfoDto(GMObjectIdentifiers.sm2sign_with_sm3.getId(), signature));
        log.info("ResponseBodySignatureAdvice.beforeBodyWrite--->end");
        return body;
    }
}
